/** 基于dbm的加密存储类
 * @file aidb.hpp
 * @author Kiven Lee
 * @version 1.0
 * @date 2022-04-08
 *
 * @copyright Kivensoft (c) 2018 - 2022
 *
 */

#include "aidb.hpp"

#include <filesystem>
#include <algorithm>

#include "klib_http/aes.h"
#include "klib_http/md5.h"

#include "klib/utils.hpp"

using namespace ai;
using namespace klib;

static constexpr char _MAGIC[4] = {'a', 'i', 'd', 'b'};
static constexpr uint8_t _VERSION[2] = {1, 0};
static constexpr size_t _AES_BLOCK_SIZE = 16;

// 生成秘钥，使用密码+向量方式生成，避免被暴力破解
static void _gen_key(uint8_t (&out)[16], string_view pass, uint8_t (&iv)[16]) {
    md5_ctx_t md5;
    md5_init(&md5);
    md5_update(&md5, pass.data(), pass.size());
    md5_update(&md5, iv, sizeof(iv));
    md5_final(&md5, out);
}

// 生成签名，使用密码+最后修改时间生成
static void _gen_sign(char (&out)[16], string_view pass, uint8_t (&salt)[2]) {
    md5_ctx_t md5;
    md5_init(&md5);
    md5_update(&md5, pass.data(), pass.size());
    md5_update(&md5, salt, sizeof(salt));
    md5_final(&md5, (uint8_t*) out);
}

// 校验签名
static bool _check_sign(string_view pass, uint8_t (&salt)[2], const char (&sign)[16]) {
    char md5_ret[16];
    _gen_sign(md5_ret, pass, salt);
    return std::equal(md5_ret, md5_ret + sizeof(md5_ret), sign);
}

// 生成随机向量
static void _fill_iv_and_salt(uint8_t (&iv)[16], uint8_t (&salt)[2]) {
    rand_t rand;
    union u64 { uint64_t v; uint8_t a[8]; } n1, n2, n3;
    n1.v = rand.next(), n2.v = rand.next(), n3.v = rand.next();
    // 随机向量
    std::copy_n(n1.a, sizeof(n1), iv);
    std::copy_n(n2.a, sizeof(n2), iv + sizeof(n1));
    // 盐值处理
    salt[0] = 0, salt[1] = 0;
    for (int i = 0; i < 4; ++i) salt[0] += n3.a[i];
    for (int i = 4; i < 8; ++i) salt[1] += n3.a[i];
}

static void _init_super_block(_aidb::_super_block_t& sb, string_view pass) {
    static_assert(sizeof(sb) % sizeof(uint64_t) == 0, "super block length check fail");
    static_assert(sizeof(_MAGIC) == sizeof(sb.magic), "magic length check fail");
    static_assert(sizeof(_VERSION) == sizeof(sb.version), "version length check fail");

    std::fill_n((uint64_t*) &sb, sizeof(sb) / sizeof(uint64_t), 0);
    std::copy_n(_MAGIC, sizeof(_MAGIC), sb.magic);
    std::copy_n(_VERSION, sizeof(_VERSION), sb.version);
    _fill_iv_and_salt(sb.iv, sb.salt);
    _gen_sign(sb.sign, pass, sb.salt);
}

aidb_value_t::aidb_value_t(string_view title, string_view user, string_view pass, string_view url, string_view memo):
        title(title), user(user), pass(pass), url(url), memo(memo) {

    size = title.size() + user.size() + pass.size() + url.size() + memo.size() + sizeof(uint32_t) * 5;
    data = new char[size];

    *((uint32_t*) data + 0) = net::h2n((uint32_t) title.size());
    *((uint32_t*) data + 1) = net::h2n((uint32_t) user.size());
    *((uint32_t*) data + 2) = net::h2n((uint32_t) pass.size());
    *((uint32_t*) data + 3) = net::h2n((uint32_t) url.size());
    *((uint32_t*) data + 4) = net::h2n((uint32_t) memo.size());

    size_t pos = sizeof(uint32_t) * 5;
    auto add_fn = [this, &pos] (string_view& s) {
        std::memcpy(this->data + pos, s.data(), s.size());
        pos += s.size();
    };

    if (title.size()) add_fn(title);
    if (user.size()) add_fn(user);
    if (pass.size()) add_fn(pass);
    if (url.size()) add_fn(url);
    if (memo.size()) add_fn(memo);
}

aidb_value_t::aidb_value_t(dbm::uni_value_t&& val): uni_value_t(std::move(val)) {
    char* d = data + sizeof(uint32_t) * 5;

    auto fn = [this, &d] (string_view& out, size_t idx) {
        if (size_t len = net::n2h(*((uint32_t*) this->data + idx)); len) {
            out = string_view(d, len);
            d += len;
        }
    };

    fn(title, 0);
    fn(user, 1);
    fn(pass, 2);
    fn(url, 3);
    fn(memo, 4);
}

aidb_t::opstate aidb_t::create(const string& filename, string_view password) {
    if (std::filesystem::exists(filename)) {
        LOG_ERROR("aidb file {} exists, create fail", filename);
        return opstate::db_exists;
    }

    // 退出函数前的资源清理工作
    bool create_ok = false;
    defer_t defer{[this, &create_ok] () { if (!create_ok) db.close(); }};

    // 打开数据库
    if (!db.open(filename.c_str(), _aidb::_AIDB_BLOCK_SIZE, open_mode_t::create)) {
        LOG_ERROR("can't create aidb file {}", filename);
        return opstate::can_not_open;
    }

    // 初始化超级块信息并写入
    _init_super_block(super_block, password);
    // 写入超级块
    if (_update_super_block() == false) {
        LOG_ERROR("write super block to aidb[{}] error", filename);
        return opstate::can_not_write;
    }

    // 生成秘钥
    _gen_key(db_key, password, super_block.iv);

    create_ok = true; // 操作成功，跳过defer函数
    return opstate::ok;
}

aidb_t::opstate aidb_t::open(const string& filename, string_view password) {
    if (!db.open(filename.c_str(), _aidb::_AIDB_BLOCK_SIZE, open_mode_t::write)) {
        LOG_ERROR("open aidb file {} fail", filename);
        return opstate::can_not_open;
    }

    // 退出函数前的资源清理工作
    bool open_ok = false;
    defer_t defer([this, &open_ok] () { if (!open_ok) db.close(); });

    // 加载超级块信息
    dbm::uni_value_t sb_val = db.fetch(aidb_key_value_t(_aidb::_SUPER_BLOCK_KEY));
    if (!sb_val) {
        LOG_ERROR("can't find super block in aidb[{}], open database fail", filename);
        return opstate::super_block_not_exists;
    }
    if (sb_val.size != sizeof(super_block)) {
        LOG_ERROR("aidb[{}] super block size[{}] error, load database fail", filename, sb_val.size);
        return opstate::super_block_size_error;
    }
    std::memcpy(&super_block, sb_val.data, sizeof(super_block));

    // 校验秘钥是否正确
    if (!_check_sign(password, super_block.salt, super_block.sign)) {
        LOG_ERROR("aidb[{}] check sign error, load database fail", filename);
        return opstate::password_error;
    }

    // 生成秘钥
    _gen_key(db_key, password, super_block.iv);

    open_ok = true; // 操作成功，跳过defer函数
    return opstate::ok;
}

aidb_value_t aidb_t::get(aidb_key_t key) {
    if (key == nkey) {
        LOG_ERROR("aidb error: can't get because key is super block key");
        return aidb_value_t();
    }

    dbm::uni_value_t val = db.fetch(aidb_key_value_t(key));
    if (val.data != nullptr) {
        assert(val.size % 16 == 0);

        char tmp[16];
        aes_ctx_t ac;
        aes_init(&ac, AES_DECRYPT, db_key, sizeof(db_key) * 8, super_block.iv);
        for (char* p = val.data, *pmax = p + val.size; p < pmax; p += 16) {
            aes_update(&ac, p, 16, tmp);
            std::copy_n(tmp, 16, p);
        }
        //截断加密字符串，去除padding字符
        val.size -= val.data[val.size - 1];
    }

    return aidb_value_t(std::move(val));
}

bool aidb_t::_put(const aidb_key_t key, const void* data, uint32_t size) {
    dbm::uni_value_t dec(PKCS7_PADDING_SIZE(size, _AES_BLOCK_SIZE));

    aes_ctx_t ac;
    aes_init(&ac, AES_ENCRYPT, db_key, sizeof(db_key) * 8, super_block.iv);
    size_t n = aes_update(&ac, data, size, dec.data);
    aes_final(&ac, dec.data + n);

    return db.store(aidb_key_value_t{key}, value_t{dec.data, dec.size}, store_mode_t::replace);
}

aidb_key_t aidb_t::firstkey() {
    aidb_key_t key_raw{};
    value_t out(&key_raw, 0);
    for (db.firstkey(out); out.size; db.nextkey(out, out)) {
        aidb_key_t key = net::n2h(key_raw);
        if (key != nkey)
            return key;
    }
    return _aidb::_SUPER_BLOCK_KEY;
}

aidb_key_t aidb_t::nextkey(aidb_key_t key) {
    aidb_key_t key_raw = net::h2n(key);
    value_t out(&key_raw, sizeof(key));
    for(db.nextkey(out, out); out.size; db.nextkey(out, out)) {
        key = net::n2h(key_raw);
        if (key != nkey)
            return key;
    }
    return _aidb::_SUPER_BLOCK_KEY;
}

#ifdef _TEST_AIDB_

using namespace klib;
using namespace ai;

int main() {
    std::string filename("1.aidb");
    std::string pass("password");
    rand_t<aidb_key_t> rand;
    aidb_key_t key;

    aidb_t db;
    if (!db.open(filename, pass) && !db.create(filename, pass)) {
        fmt::println("打开数据库{}失败", filename);
        return -1;
    }

    klib::dstring_t ds;
    for (int i = 0; i < 10; i++) {
        key = rand.next();
        ds.resize(0);
        ds.format("value is {}", key);
        db.put(ds);
    }

    for (;;) {
        fmt::println("================================================================================");
        int index = 100;
        for (auto k : db) {
            auto v = db.get(k);
            fmt::println("{}: key = {}, value = {}", ++index, k, std::string_view(v.data, v.size));
        }
        fmt::println("================================================================================");

        std::printf("输出要删除的键： ");
        std::scanf("%u", &key);
        if (key == 0) break;
        auto v = db.get(key);
        if (v.data) {
            fmt::println("检索结果： key = {}, value = {}", key, std::string_view(v.data, v.size));
            db.erase(key);
        } else {
            fmt::println("未找到记录：{}", key);
        }
    }
    return 0;
}

#endif // _TEST_AIDB_